package org.example;

import java.io.IOException;
import java.net.InetAddress;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.*;
import java.util.Arrays;
import java.util.Iterator;
import java.util.Queue;
import java.util.concurrent.ConcurrentLinkedQueue;


public class OptMain {

    public static void main(String[] args) throws IOException {

        //创建一个tcp服务端
        int port = 1080;
        Selector selector = Selector.open();
        ServerSocketChannel serverSocketChannel = ServerSocketChannel.open();
        serverSocketChannel.bind(new InetSocketAddress(port));
        serverSocketChannel.configureBlocking(false);
        serverSocketChannel.register(selector, SelectionKey.OP_ACCEPT);

        System.out.println("tcp服务端启动成功,端口："+port);


        while (!Thread.interrupted()) {

            selector.select(200);

            Iterator<SelectionKey> iter = selector.selectedKeys().iterator();
            while (iter.hasNext()) {
                SelectionKey key = iter.next();
                iter.remove();
                try{

                    if (!key.isValid()) {
                        continue; // 如果键无效，跳过处理
                    }

                    if (key.isAcceptable()) {
                        //服务端接收客户端请求
                        handleAccept(serverSocketChannel, selector);
                    } else if(key.isConnectable()){
                        //本地连接
                        handleConnect(key, selector);
                    }else{
                        SocketChannel socketChannel = (SocketChannel) key.channel();
                        ChannelContext context = (ChannelContext)key.attachment();
                        if (key.isValid() && key.isReadable()) {
                            handleRead(socketChannel, context, key, selector);
                        }

                        if (key.isValid() && key.isWritable()) {
                            handleWrite(context, socketChannel, key);
                        }

                    }

                }catch (Exception e){
                    e.printStackTrace();
                }

            }

        }
    }

    private static void handleWrite(ChannelContext context, SocketChannel socketChannel, SelectionKey key) {
        try {
            Queue<ByteBuffer> queue = context.selfWriteQueue;
            ByteBuffer peek = queue.peek();

            while (peek != null) {
                int written = socketChannel.write(peek);
                if (written == 0) {
                    // 写不动了，等下次写事件
                    break;
                }
                if (!peek.hasRemaining()) {
                    queue.poll(); // 移除已写完的 buffer
                    peek = queue.peek(); // 尝试下一个
                } else {
                    // 本次没写完，退出等下次继续
                    break;
                }
            }

            // 如果写完了所有数据，取消写事件监听
            if (queue.isEmpty()) {
                key.interestOps(key.interestOps() & ~SelectionKey.OP_WRITE);
            }

        } catch (IOException e) {
            e.printStackTrace();
            key.cancel(); // 取消 key
            try {
                socketChannel.close();
            } catch (IOException ex) {
                ex.printStackTrace();
            }
        }
    }

    private static void handleRead(SocketChannel socketChannel, ChannelContext context, SelectionKey key, Selector selector) throws IOException {
        ByteBuffer buffer = ByteBuffer.allocate(1024);
        try{
            int bytesRead = socketChannel.read(buffer);
            if (bytesRead == -1) {
                socketChannel.close();
                if(context.dstKey != null) {
                    context.dstKey.channel().close();
                }
            } else if (bytesRead > 0) {
                buffer.flip();

                if(context.state == State.NEGOTIATING) {
                    //第一步 协商连接
                    negotiatedConnection(key, buffer);
                    context.state = State.CONNECTING;
                }else if(context.state == State.CONNECTING){
                    //第二步 连接客户端发送的ip和端口
                    connectServer(selector, key,buffer, context);
                    context.state = State.FORWARD;
                }else{
                    //一个channel写入另一个channel
                    SelectionKey dstKey = context.dstKey;
                    nioWrite(dstKey, buffer);
                }
            }
        }catch (Exception e){
            e.printStackTrace();
            socketChannel.close();
            if(context.dstKey != null) {
                context.dstKey.channel().close();
            }
        }
    }

    private static void handleConnect(SelectionKey key, Selector selector) throws IOException {
        SocketChannel socketChannel = (SocketChannel) key.channel();

        ChannelContext context = (ChannelContext) key.attachment();
        SelectionKey dstKey =context.dstKey;
        byte[] replyData = context.replyData;
        SocketChannel dstChannel = (SocketChannel)dstKey.channel();

        //目标上下文
        ChannelContext dstContext = (ChannelContext)dstKey.attachment();

        try {
            // 完成连接[这里会出现超时异常]
            socketChannel.finishConnect();

            socketChannel.register(selector,SelectionKey.OP_READ,context);

            //目标是一直在读的状态，只要设置目标key
            dstContext.dstKey = key;

            //告诉客户端连接成功
            nioWrite(dstKey, ByteBuffer.wrap(replyData));

        } catch (Exception ex) {
            ex.printStackTrace();
            socketChannel.close();
            dstChannel.close();
        }
    }

    private static void handleAccept(ServerSocketChannel serverSocketChannel, Selector selector) throws IOException {
        SocketChannel socketChannel = serverSocketChannel.accept();
        socketChannel.configureBlocking(false);
        ChannelContext context = new ChannelContext(State.NEGOTIATING);
        socketChannel.register(selector, SelectionKey.OP_READ,context);
    }

    enum State {
        NONE("空"),NEGOTIATING("协商连接"), CONNECTING("连接服务"), FORWARD("数据转发");
        State(String desc) {
        }
    }

    public static class ChannelContext{
        public State state;
        public SelectionKey dstKey; //目标key
        public Queue<ByteBuffer> selfWriteQueue = new ConcurrentLinkedQueue<>();
        public byte[] replyData; //回复数据

        public ChannelContext(State state) {
            this.state = state;
        }
    }

    private static void nioWrite(SelectionKey dstKey, ByteBuffer buffer) throws IOException {

        SocketChannel socketChannel = (SocketChannel) dstKey.channel();
        ChannelContext context = (ChannelContext) dstKey.attachment();
        Queue<ByteBuffer> queue = context.selfWriteQueue;

        // 如果队列为空，说明可以尝试立即写
        if (queue.isEmpty()) {
            socketChannel.write(buffer);
            if (!buffer.hasRemaining()) {
                return; // 全部写完了，不需要入队
            }
        }

        // 仍有剩余或前面已有 buffer，加入队列
        ByteBuffer toQueue = buffer.duplicate();
        queue.add(toQueue);

        // 注册写事件
        dstKey.interestOps(dstKey.interestOps() | SelectionKey.OP_WRITE);

    }


    private static void negotiatedConnection(SelectionKey key,ByteBuffer buffer) throws IOException {
        //SOCKS版本
        byte version = buffer.get();
        //认证方法数量
        byte nmethods = buffer.get();
        //认证方法
        boolean noAuth = false;
        byte[] methods = new byte[nmethods];
        for (byte b = 0; b < nmethods; b++) {
            methods[b] = buffer.get();
            //无需验证
            if(methods[b] == 0){
                noAuth = true;
            }
        }

        //协议是socks5，支持无需要验证
        if(version == 5 && noAuth){
            //回复客户端,版本5，无需验证
            byte[] reply = {5,0};
            nioWrite(key, ByteBuffer.wrap(reply));
        }else{
            //不支持关闭连接
            key.channel().close();
        }
    }
    private static void connectServer(Selector selector, SelectionKey key, ByteBuffer buffer, ChannelContext context) throws IOException {
        //SOCKS版本
        byte version = buffer.get();
        //CMD：命令，1 表示连接，2 表示绑定。
        byte cmd = buffer.get();
        //保留字段，必须为 0x00。
        byte rsv = buffer.get();
        //地址类型，1 表示 IPv4 地址，3 表示域名，4 表示 IPv6 地址
        byte atyp = buffer.get();
        String host = "";
        if(atyp == 1){
            //ipv4
            byte[] bt = new byte[4];
            buffer.get(bt);
            host = InetAddress.getByAddress(bt).getHostAddress();

        }else if(atyp == 3){
            //域名
            byte len = buffer.get();
            byte[] bt = new byte[len];
            buffer.get(bt);
            host = new String(bt);

        }else if(atyp == 4){
            //ipv6
            byte[] bt = new byte[16];
            buffer.get(bt);
            host = InetAddress.getByAddress(bt).getHostAddress();
        }

        int port = Short.toUnsignedInt(buffer.getShort());

        System.out.println("host:"+host+",port:"+port);

        InetSocketAddress address = new InetSocketAddress(host, port);
        SocketChannel clientChannel = SocketChannel.open();
        clientChannel.configureBlocking(false);
        clientChannel.connect(address);

        byte[] replyData = Arrays.copyOf(buffer.array(), buffer.position());
        //应答成功
        replyData[1] = 0;

        ChannelContext context1 = new ChannelContext(State.NONE);
        context1.replyData = replyData;
        context1.dstKey = key;
        clientChannel.register(selector,SelectionKey.OP_CONNECT,context1);

    }
}
